"""数据清洗与目标构建工具."""

from __future__ import annotations

from typing import Tuple

import pandas as pd

from .schemas import ETFHistoricalDataset


def prepare_training_dataset(
    features_df: pd.DataFrame,
    dataset: ETFHistoricalDataset,
    *,
    target_column: str = "target_up",
    min_samples: int = 60,
) -> Tuple[pd.DataFrame, pd.Series]:
    """基于缓存特征与历史行情生成可训练的数据集.

    参数
    ----
    features_df:
        特征工程后的数据框，行索引需与行情数据对齐。
    dataset:
        `ETFHistoricalDataset` 实例，包含原始价格数据。
    target_column:
        特征数据中目标列的名称，缺省为 ``target_up``。
    min_samples:
        清洗后允许的最小样本数，默认 60。

    返回
    ----
    Tuple[pd.DataFrame, pd.Series]
        清洗后的特征矩阵和二分类目标序列。

    Raises
    ------
    ValueError
        当输入为空、缺少必要列或样本数不足时抛出异常。
    """

    if features_df is None or features_df.empty:
        raise ValueError("features_df 不能为空，且需要包含特征列。")
    if dataset is None or dataset.frame.empty:
        raise ValueError("dataset 不能为空，且需要包含历史行情数据。")

    # 优先使用特征工程阶段生成的标签，否则就地构造“下一日是否上涨”的目标。
    if target_column in features_df.columns:
        y = features_df[target_column]
        X = features_df.drop(columns=[target_column])
    else:
        if "收盘" not in dataset.frame.columns:
            raise ValueError("行情数据缺少“收盘”列，无法构造目标变量。")
        close = dataset.frame["收盘"]
        # 生成二分类标签：下一交易日收盘价高于当日则记为 1，否则为 0。
        y = (close.shift(-1) > close).astype(int)
        try:
            y = y.loc[features_df.index]
        except KeyError as exc:
            raise ValueError("特征索引与行情索引不一致，无法对齐目标变量。") from exc
        X = features_df

    # 去除包含缺失值的样本，避免模型训练过程中产生 NaN。
    valid_mask = ~(X.isna().any(axis=1) | y.isna())
    X_clean = X[valid_mask]
    y_clean = y[valid_mask]

    if len(X_clean) < min_samples:
        raise ValueError(
            f"数据量不足：有效样本数 {len(X_clean)} 条，至少需要 {min_samples} 条才能训练模型。"
        )

    return X_clean, y_clean

